package daybook.util;

import org.json.JSONArray;

import java.util.*;
import java.util.function.*;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

public final class StreamUtils {
    private StreamUtils() {}

    public static <T, R> Iterator<R> mapWithIndex(Iterator<T> iterator, BiFunction<T, Integer, R> mapper) {
        return new Iterator<R>() {
            private int index = 0;
            @Override
            public boolean hasNext() {
                return iterator.hasNext();
            }

            @Override
            public R next() {
                return mapper.apply(iterator.next(), index++);
            }
        };
    }

    public static <T, A, R> R collectIterator(Iterator<T> iterator, Collector<T, A, R> collector) {
        final BiConsumer<A, T> accumulator = collector.accumulator();
        final A result = collector.supplier().get();
        while (iterator.hasNext()) {
            accumulator.accept(result, iterator.next());
        }
        return collector.finisher().apply(result);
    }

    public static <T, A> A reduceIterator(Iterator<T> iterator, A zero, BiFunction<A, T, A> accumulator) {
        A result = zero;
        while (iterator.hasNext()) {
            result = accumulator.apply(result, iterator.next());
        }
        return result;
    }

    public static <T> T reduceIterator(Iterator<T> iterator, T zero, BinaryOperator<T> combiner) {
        T result = zero;
        while (iterator.hasNext()) {
            result = combiner.apply(result, iterator.next());
        }
        return result;
    }

    public static <T, A> Iterator<A> scanIterator(Iterator<T> iterator, A zero, BiFunction<A, T, A> accumulator) {
        return new Iterator<A>() {
            private boolean started = false;
            private A elem;

            @Override
            public boolean hasNext() {
                return !started || iterator.hasNext();
            }

            @Override
            public A next() {
                if (started) {
                    elem = accumulator.apply(elem, iterator.next());
                } else {
                    elem = zero;
                    started = true;
                }
                return elem;
            }
        };
    }

    public static <T, A> Stream<A> scanStream(Stream<T> stream, A zero, BiFunction<A, T, A> accumulator) {
        return StreamSupport.stream(Spliterators.spliteratorUnknownSize(
                scanIterator(stream.iterator(), zero, accumulator),
                Spliterator.ORDERED | Spliterator.IMMUTABLE), false);
    }

    public static <T, U, A, R>
    Collector<T, ?, R> flatMapping(Function<? super T, ? extends Stream<? extends U>> mapper,
                               Collector<? super U, A, R> downstream) {
        final BiConsumer<A, ? super U> downstreamAccumulator = downstream.accumulator();
        final BiConsumer<A, T> accumulator = (a, t) -> {
            mapper.apply(t).forEach(u -> {
                downstreamAccumulator.accept(a, u);
            });
        };

        if (downstream.characteristics().contains(Collector.Characteristics.IDENTITY_FINISH)) {
            @SuppressWarnings("unchecked")
            final Collector<? super U, R, R> ds = (Collector<? super  U, R, R>) downstream;
            @SuppressWarnings("unchecked")
            final BiConsumer<R, T> acc = (BiConsumer<R, T>) accumulator;
            return Collector.of(ds.supplier(), acc, ds.combiner());
        } else {
            return Collector.of(
                    downstream.supplier(),
                    accumulator,
                    downstream.combiner(),
                    downstream.finisher());
        }
    }

    @SuppressWarnings("unchecked")
    private static <I, R> Function<I, R> castingIdentity() {
        return i -> (R) i;
    }

    public static <T, K, D, A>
    Collector<T, ?, Map<K, D>> multiGroupingBy(Function<? super T, ? extends Iterable<? extends K>> classifier,
                                       Collector<? super T, A, D> downstream) {
        return multiGroupingBy(classifier, HashMap::new, downstream);
    }

    public static <T, K, D, A, M extends Map<K, D>>
    Collector<T, ?, M> multiGroupingBy(Function<? super T, ? extends Iterable<? extends K>> classifier,
                                  Supplier<M> mapFactory,
                                  Collector<? super T, A, D> downstream) {
        final Supplier<A> downstreamSupplier = downstream.supplier();
        final BiConsumer<A, ? super T> downstreamAccumulator = downstream.accumulator();
        final BiConsumer<Map<K, A>, T> accumulator = (m, t) -> {
            for (K key: classifier.apply(t)) {
                Objects.requireNonNull(key, "element cannot be mapped to a null key");
                final A container = m.computeIfAbsent(key, k -> downstreamSupplier.get());
                downstreamAccumulator.accept(container, t);
            }
        };
        final BinaryOperator<Map<K, A>> merger = (m1, m2) -> Utils.merge(m1, m2, downstream.combiner());
        @SuppressWarnings("unchecked")
        final Supplier<Map<K, A>> mangledFactory = (Supplier<Map<K, A>>) mapFactory;

        if (downstream.characteristics().contains(Collector.Characteristics.IDENTITY_FINISH)) {
            return Collector.of(mangledFactory, accumulator, merger, castingIdentity(), Collector.Characteristics.IDENTITY_FINISH);
        } else {
            @SuppressWarnings("unchecked")
            final Function<A, A> downstreamFinisher = (Function<A, A>) downstream.finisher();
            final Function<Map<K, A>, M> finisher = intermediate -> {
                intermediate.replaceAll((k, v) -> downstreamFinisher.apply(v));
                @SuppressWarnings("unchecked")
                final M castResult = (M) intermediate;
                return castResult;
            };
            return Collector.of(mangledFactory, accumulator, merger, finisher);
        }
    }

    public static <T, A, R> Collector<T, ?, R> filter(Predicate<? super T> predicate, Collector<? super T, A, R> downstream) {
        final BiConsumer<A, ? super T> downstreamAccumulator = downstream.accumulator();
        final BiConsumer<A, T> accumulator = (a, t) -> {
            if (predicate.test(t)) {
                downstreamAccumulator.accept(a, t);
            }
        };

        return Collector.of(
                downstream.supplier(),
                accumulator,
                downstream.combiner(),
                downstream.finisher(),
                downstream.characteristics().toArray(new Collector.Characteristics[0])
        );
    }

    private static Collector<Object, JSONArray, JSONArray> JSON_ARRAY_COLLECTOR = Collector.of(
            JSONArray::new, JSONArray::put, JSONArray::put
    );
    public static Collector<Object, JSONArray, JSONArray> jsonArrayCollector() {
        return JSON_ARRAY_COLLECTOR;
    }

    private static <T> BinaryOperator<T> throwingMerger() {
        return (u,v) -> { throw new IllegalStateException(String.format("Duplicate key %s", u)); };
    }

    public static <T, K, U, M extends Map<K, U>>
    Collector<T, ?, M> toMap(Function<? super T, ? extends K> keyMapper,
                             Function<? super T, ? extends U> valueMapper,
                             Supplier<M> mapSupplier) {
        return Collectors.toMap(keyMapper, valueMapper, throwingMerger(), mapSupplier);
    }
}
